import cv2
import numpy as np
from multiprocessing import Queue
from typing import Any, Dict, List, Tuple
import time
import yaml


def saveCalibration(filename: str, data: Dict[str, Any]):
    """保存相机标定结果到文件"""
    with open(filename, "w") as f:
        yaml.dump(data, f)


def loadCalibration(filename: str) -> Dict[str, Any]:
    """从文件加载相机标定结果"""
    with open(filename) as f:
        data = yaml.load(f, Loader=yaml.FullLoader)
    for key in ["camera_matrix", "dist_coeff", "R", "T", "E", "F"]:
        if key in data:
            data[key] = np.array(data[key])
    return data


def captureCalibrationImages(
    shared_state: Dict[str, Any], num_images=20
) -> Tuple[List[str], List[str], Tuple[int, int]]:
    """采集标定图像

    Args:
        shared_state (dict): 包含相机队列的共享状态
        num_images (int): 需要采集的图像数量

    Returns:
        list: 左相机采集到的图像路径
        list: 右相机采集到的图像路径
        tuple: 图像的尺寸 (宽, 高)
    """
    left_camera_queue = shared_state["left_camera_queue"]
    right_camera_queue = shared_state["right_camera_queue"]

    left_captured_images = []
    right_captured_images = []

    print(
        f"请在终端中按照提示采集 {num_images} 张棋盘格图像. 按空格键拍摄，按 'q' 退出。"
    )

    image_size = None

    while (
        len(left_captured_images) < num_images
        or len(right_captured_images) < num_images
    ):
        while left_camera_queue.empty() or right_camera_queue.empty():
            time.sleep(0.1)

        left_frame = left_camera_queue.get()
        right_frame = right_camera_queue.get()

        combined_frame = np.hstack((left_frame, right_frame))
        cv2.imshow("frame", combined_frame)
        key = cv2.waitKey(1) & 0xFF

        if key == ord(" "):  # 按空格键拍摄
            if image_size is None:
                image_size = (left_frame.shape[1], left_frame.shape[0])

            left_image_path = (
                f"./images/calib_left_image_{len(left_captured_images) + 1}.png"
            )
            right_image_path = (
                f"./images/calib_right_image_{len(right_captured_images) + 1}.png"
            )
            cv2.imwrite(left_image_path, left_frame)
            cv2.imwrite(right_image_path, right_frame)

            left_captured_images.append(left_image_path)
            right_captured_images.append(right_image_path)
            print(f"已采集 {len(left_captured_images)} 对图像")
        elif key == ord("q"):  # 按 'q' 键退出
            break

    cv2.destroyAllWindows()

    return left_captured_images, right_captured_images, image_size


def calibrateCamera(
    image_paths: List[str],
    chessboard_size: Tuple[int, int] = (9, 6),
    square_size: float = 1.0,
):
    """相机标定

    Args:
        image_paths (list): 采集的标定图像路径列表
        chessboard_size (tuple): 棋盘格内角点的行列数
        square_size (float): 单个棋盘格的边长 cm

    Returns:
        ret (float): 标定误差
        mtx (np.ndarray): 相机内参矩阵
        dist (np.ndarray): 相机畸变系数
        rvecs (list): 旋转向量
        tvecs (list): 平移向量
        objpoints (list): 对象点列表
        imgpoints (list): 图像点列表
    """
    # 终止条件
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 1e-2)

    # 内角点数量
    cols = chessboard_size[0] - 1
    rows = chessboard_size[1] - 1

    # 准备对象点, 如 (0,0,0), (1,0,0), (2,0,0) ....,(cols,rows)
    objp = np.zeros((cols * rows, 3), np.float32)
    objp[:, :2] = np.mgrid[0:cols, 0:rows].T.reshape(-1, 2)
    objp *= square_size

    # 存储所有图像的对象点和图像点
    objpoints = []  # 3d 点在真实世界空间
    imgpoints = []  # 2d 点在图像平面

    for fname in image_paths:
        img = cv2.imread(fname)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        # 查找棋盘格角点
        ret, corners = cv2.findChessboardCorners(gray, (cols, rows), None)

        # 如果找到，添加对象点和图像点
        if ret:
            print(f"在图像 {fname} 中找到 {len(corners)} 个角点")
            if len(corners) == cols * rows:
                objpoints.append(objp)
                corners_subpix = cv2.cornerSubPix(
                    gray, corners, (11, 11), (-1, -1), criteria
                )

                imgpoints.append(corners_subpix)

                # 绘制并显示角点
                img = cv2.drawChessboardCorners(img, (cols, rows), corners_subpix, ret)
                cv2.imshow("Corners", img)
                cv2.waitKey(500)  # 显示图像 0.5 秒

            else:
                print(f"警告: 图像 {fname} 中的角点数量不正确，可能需要检查棋盘格尺寸")
        else:
            print(f"未能在图像 {fname} 中找到棋盘格角点")

    cv2.destroyAllWindows()

    if not objpoints or not imgpoints:
        print("没有检测到足够的角点来执行相机标定")
        return None, None, None, None, None, None, None

    # 标定相机
    ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(
        objpoints, imgpoints, gray.shape[::-1], None, None
    )

    return ret, mtx, dist, rvecs, tvecs, objpoints, imgpoints


def runCalibration(
    chessboard_size: Tuple[int, int],
    num_images: int,
    left_image_paths: List[str],
    right_image_paths: List[str],
) -> None:
    """执行相机校准

    Args:
        chessboard_size (Tuple[int, int]): 棋盘格的行列数
        num_images (int): 需要采集的图像数量
        left_image_paths (List[str]): 左相机采集的图像路径列表
        right_image_paths (List[str]): 右相机采集的图像路径列表
    """
    print("开始采集标定图像...")

    if len(left_image_paths) < num_images or len(right_image_paths) < num_images:
        print(
            f"采集的图像数量不足. 左相机图像数量: {len(left_image_paths)}, 右相机图像数量: {len(right_image_paths)}"
        )
        return

    print("开始左相机标定...")
    left_ret, left_mtx, left_dist, left_rvecs, left_tvecs, _, _ = calibrateCamera(
        left_image_paths, chessboard_size
    )
    if left_ret is not None:
        print("左相机标定成功")
        print("左相机内参矩阵:\n", left_mtx)
        print("左相机畸变系数:\n", left_dist)
    else:
        print("左相机标定失败")
        left_mtx, left_dist = None, None

    print("开始右相机标定...")
    right_ret, right_mtx, right_dist, right_rvecs, right_tvecs, _, _ = calibrateCamera(
        right_image_paths, chessboard_size
    )
    if right_ret is not None:
        print("右相机标定成功")
        print("右相机内参矩阵:\n", right_mtx)
        print("右相机畸变系数:\n", right_dist)
    else:
        print("右相机标定失败")
        right_mtx, right_dist = None, None

        # 保存标定结果到文件
        saveCalibration(
            "./config/left_camera_calibration.yaml",
            {"camera_matrix": left_mtx, "dist_coeff": left_dist},
        )
        saveCalibration(
            "./config/right_camera_calibration.yaml",
            {"camera_matrix": right_mtx, "dist_coeff": right_dist},
        )

        print(
            "标定完成，结果已保存到 config 目录下的 left_camera_calibration.yaml 和 right_camera_calibration.yaml 文件中"
        )


def runStereoCalibration(
    chessboard_size: Tuple[int, int],
    num_images: int,
    left_image_paths: List[str],
    right_image_paths: List[str],
    image_size: Tuple[int, int],
) -> None:
    """执行双目相机校准

    Args:
        chessboard_size (Tuple[int, int]): 棋盘格的行列数
        num_images (int): 需要采集的图像数量
        left_image_paths (List[str]): 左相机采集的图像路径列表
        right_image_paths (List[str]): 右相机采集的图像路径列表
        image_size (Tuple[int, int]): 图像的尺寸 (宽, 高)
    """

    size = image_size
    if size is None:
        img = cv2.imread(left_image_paths[0])
        size = (img.shape[1], img.shape[0])

    square_size = 7.0 / 9

    if len(left_image_paths) < num_images or len(right_image_paths) < num_images:
        print(
            f"采集的图像数量不足. 左相机图像数量: {len(left_image_paths)}, 右相机图像数量: {len(right_image_paths)}"
        )
        return

    print("开始左相机标定...")
    (
        left_ret,
        left_mtx,
        left_dist,
        left_rvecs,
        left_tvecs,
        left_objpoints,
        left_imgpoints,
    ) = calibrateCamera(left_image_paths, chessboard_size, square_size)
    if left_ret is not None:
        print("左相机标定成功")
    else:
        print("左相机标定失败")
        return

    print("开始右相机标定...")
    (
        right_ret,
        right_mtx,
        right_dist,
        right_rvecs,
        right_tvecs,
        right_objpoints,
        right_imgpoints,
    ) = calibrateCamera(right_image_paths, chessboard_size, square_size)
    if right_ret is not None:
        print("右相机标定成功")
    else:
        print("右相机标定失败")
        return

    """使用一组对象点（棋盘格的 3D 坐标）和图像点（左右相机中的 2D 像素坐标）
    计算双目相机的内参、畸变系数、旋转矩阵、平移向量、本质矩阵和基础矩阵
        - retval: 重投影误差
        - camera_matrix_left: 校正后的左相机内参矩阵
        - dist_coeffs_left: 校正后的左相机畸变系数
        - camera_matrix_right: 校正后的右相机内参矩阵
        - dist_coeffs_right: 校正后的右相机畸变系数
        - R: 两个相机之间的旋转矩阵
        - T: 两个相机之间的平移向量
        - E: 本质矩阵，反映两个相机之间的几何关系
        - F: 基础矩阵，用于将一个相机的图像点映射到另一个相机的图像平面上
    """
    print("开始双目相机标定...")
    (
        retval,
        camera_matrix_left,
        dist_coeffs_left,
        camera_matrix_right,
        dist_coeffs_right,
        R,
        T,
        E,
        F,
    ) = cv2.stereoCalibrate(
        objectPoints=left_objpoints,
        imagePoints1=left_imgpoints,
        imagePoints2=right_imgpoints,
        cameraMatrix1=left_mtx,
        distCoeffs1=left_dist,
        cameraMatrix2=right_mtx,
        distCoeffs2=right_dist,
        imageSize=size,
        criteria=(cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 1e-5),
        flags=cv2.CALIB_FIX_INTRINSIC,
    )

    print("双目相机标定完成")
    print("左相机内参矩阵:\n", camera_matrix_left)
    print("左相机畸变系数:\n", dist_coeffs_left)
    print("右相机内参矩阵:\n", camera_matrix_right)
    print("右相机畸变系数:\n", dist_coeffs_right)
    print("旋转矩阵:\n", R)
    print("平移向量:\n", T)
    print("本质矩阵:\n", E)
    print("基础矩阵:\n", F)

    # 保存标定结果到文件
    stereo_data = {
        "left_camera_matrix": camera_matrix_left.tolist(),
        "left_dist_coeffs": dist_coeffs_left.tolist(),
        "right_camera_matrix": camera_matrix_right.tolist(),
        "right_dist_coeffs": dist_coeffs_right.tolist(),
        "R": R.tolist(),
        "T": T.tolist(),
        "E": E.tolist(),
        "F": F.tolist(),
    }
    saveCalibration("./config/stereo_calibration.yaml", stereo_data)
    print("标定结果已保存到文件 './config/stereo_calibration.yaml'")

    print("\n计算矫正映射...")
    R1, R2, P1, P2, Q, validPixROI1, validPixROI2 = cv2.stereoRectify(
        cameraMatrix1=camera_matrix_left,
        distCoeffs1=dist_coeffs_left,
        cameraMatrix2=camera_matrix_right,
        distCoeffs2=dist_coeffs_right,
        imageSize=size,
        R=R,
        T=T,
        flags=cv2.CALIB_ZERO_DISPARITY,
        # alpha=0,
    )
    print(f"左相机旋转矩阵R1:\n{R1}")
    print(f"右相机旋转矩阵R2:\n{R2}")
    print(f"左相机投影矩阵P1:\n{P1}")
    print(f"右相机投影矩阵P2:\n{P2}")
    print(f"重投影矩阵Q:\n{Q}")

    left_map1, left_map2 = cv2.initUndistortRectifyMap(
        cameraMatrix=camera_matrix_left,
        distCoeffs=dist_coeffs_left,
        R=R1,
        newCameraMatrix=P1,
        size=size,
        m1type=cv2.CV_16SC2,
    )
    right_map1, right_map2 = cv2.initUndistortRectifyMap(
        cameraMatrix=camera_matrix_right,
        distCoeffs=dist_coeffs_right,
        R=R2,
        newCameraMatrix=P2,
        size=size,
        m1type=cv2.CV_16SC2,
    )

    print("完成矫正映射计算")

    """
        - left_map1: 校正后的图像的x坐标映射
        - left_map2: 校正后的图像的y坐标映射
        - right_map1: 同上
        - right_map2: 同上
    """
    # 保存矫正映射矩阵
    np.save("./config/Q.npy", Q)
    np.save("./config/left_map1.npy", left_map1)
    np.save("./config/left_map2.npy", left_map2)
    np.save("./config/right_map1.npy", right_map1)
    np.save("./config/right_map2.npy", right_map2)
    print(
        "矫正映射矩阵已保存到文件 './config/Q.npy', './config/left_map1.npy', './config/left_map2.npy', './config/right_map1.npy', './config/right_map2.npy'"
    )
